#!/bin/bash

# EnPlatform Assistant 模型管理脚本
# 自动下载和配置Excel处理优化的AI模型

set -e

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m'

log_info() {
    echo -e "${BLUE}[MODEL-MANAGER]${NC} $1"
}

log_success() {
    echo -e "${GREEN}[MODEL-MANAGER]${NC} $1"
}

log_warning() {
    echo -e "${YELLOW}[MODEL-MANAGER]${NC} $1"
}

log_error() {
    echo -e "${RED}[MODEL-MANAGER]${NC} $1"
}

# 等待服务就绪
wait_for_service() {
    local service_name=$1
    local service_url=$2
    local max_attempts=30
    local attempt=1

    log_info "等待 $service_name 服务就绪..."
    
    while [ $attempt -le $max_attempts ]; do
        if curl -s "$service_url/api/tags" >/dev/null 2>&1; then
            log_success "$service_name 服务已就绪"
            return 0
        fi
        
        log_info "等待 $service_name 服务启动... ($attempt/$max_attempts)"
        sleep 10
        attempt=$((attempt + 1))
    done
    
    log_error "$service_name 服务启动超时"
    return 1
}

# 下载模型
download_model() {
    local service_name=$1
    local service_url=$2
    local model_name=$3
    local description=$4
    
    log_info "在 $service_name 中下载模型: $model_name ($description)"
    
    # 检查模型是否已存在
    if curl -s "$service_url/api/tags" | grep -q "$model_name"; then
        log_success "模型 $model_name 已存在，跳过下载"
        return 0
    fi
    
    # 下载模型
    local container_name=""
    case $service_name in
        "Ollama主服务")
            container_name="enplatform-ollama"
            ;;
        "英文嵌入服务")
            container_name="enplatform-english-embeddings"
            ;;
        "中文嵌入服务")
            container_name="enplatform-chinese-embeddings"
            ;;
    esac
    
    if [ -n "$container_name" ]; then
        log_info "开始下载 $model_name..."
        if docker exec "$container_name" ollama pull "$model_name"; then
            log_success "模型 $model_name 下载成功"
        else
            log_error "模型 $model_name 下载失败"
            return 1
        fi
    else
        log_error "未知服务名称: $service_name"
        return 1
    fi
}

# 主要模型下载流程
main() {
    log_info "开始Excel处理优化模型下载流程..."
    
    # 等待所有服务就绪
    wait_for_service "Ollama主服务" "http://ollama:11434" || exit 1
    wait_for_service "英文嵌入服务" "http://english-embeddings:11434" || exit 1
    
    # 检查是否有中文嵌入服务
    if curl -s "http://chinese-embeddings:11434/api/tags" >/dev/null 2>&1; then
        wait_for_service "中文嵌入服务" "http://chinese-embeddings:11434" || exit 1
        CHINESE_SERVICE_AVAILABLE=true
    else
        log_warning "中文嵌入服务不可用，跳过中文模型下载"
        CHINESE_SERVICE_AVAILABLE=false
    fi
    
    # 下载对话模型
    log_info "=== 下载对话模型 ==="
    download_model "Ollama主服务" "http://ollama:11434" "qwen2.5:7b" "Excel数据分析专用对话模型" || {
        log_warning "主要对话模型下载失败，尝试备用模型..."
        download_model "Ollama主服务" "http://ollama:11434" "phi3.5:3.8b" "轻量级备用对话模型"
    }
    
    # 下载英文嵌入模型
    log_info "=== 下载英文嵌入模型 ==="
    download_model "英文嵌入服务" "http://english-embeddings:11434" "nomic-embed-text" "英文和技术术语嵌入模型"
    
    # 下载中文嵌入模型（如果服务可用）
    if [ "$CHINESE_SERVICE_AVAILABLE" = true ]; then
        log_info "=== 下载中文嵌入模型 ==="
        download_model "中文嵌入服务" "http://chinese-embeddings:11434" "bge-m3" "中文和多语言嵌入模型" || {
            log_warning "主要中文模型下载失败，尝试备用模型..."
            download_model "中文嵌入服务" "http://chinese-embeddings:11434" "bge-small-zh-v1.5" "轻量级中文嵌入模型"
        }
    fi
    
    # 验证模型下载
    log_info "=== 验证模型下载 ==="
    
    log_info "Ollama主服务模型列表:"
    docker exec enplatform-ollama ollama list || log_warning "无法获取主服务模型列表"
    
    log_info "英文嵌入服务模型列表:"
    docker exec enplatform-english-embeddings ollama list || log_warning "无法获取英文嵌入模型列表"
    
    if [ "$CHINESE_SERVICE_AVAILABLE" = true ]; then
        log_info "中文嵌入服务模型列表:"
        docker exec enplatform-chinese-embeddings ollama list || log_warning "无法获取中文嵌入模型列表"
    fi
    
    log_success "模型下载流程完成！"
    log_info "现在可以使用以下功能："
    log_info "  - Excel文件解析和嵌入"
    log_info "  - 基于Excel数据的对话问答"
    log_info "  - 多语言语义搜索"
    
    # 创建完成标记文件
    touch /tmp/model-download-complete
}

# 执行主流程
main "$@"
